from functionSim_config import *
import torch
from torch.utils.data import Dataset
import shelve
import numpy as np
from tqdm import tqdm
import copy
from easySample import easySample
from dataPre import dataPre
import sys
sys.path.append(r"/home/cyw/projects/function_sim_project/basic_script")


# load_from_shelve=True
load_from_shelve = False


class functionSimDataset(Dataset):
    """
        这里将样本调调整成适合functionSim模型使用的格式
        主要的功能，将样本转为batch  batch支持1
        输入格式为：
            一对样本名和value  可能存在很多
        样本的输出格式为：
            样本应该为一对对的，每一对是两个，为a,b,value
        能区分batch,支持batch为1的情况
        其他情况，使用get_samples就好，在区分batch

        8/15大更新,这个函数变成dataloader,内存不足狠狠的给我上了一课
        8/16大更新，输入变成a,b,c了，但是考虑到a,b的情况，所以应该是四维value，a,b,c
            训练时只用a,b和b,c，用模型得出相似值时用a,b即可，c留空，value用于表示后者情况的真实值
    """

    def __init__(self, datas, datasName, is_train):
        """
            datas   [(x,y,value),(x,y,value),(x,y,value)]
            x,y为样本的md5值
            初始化是将需要用的样本加载进内存    
        """
        self.data = datas
        self.samples = {}
        self.eSample = easySample()
        self.is_train = is_train
        self.datasName = datasName
        self.tripletDatas = []

        # print("{}正在加载样本,样本长度：{}".format(modelName,len(datas)))
        # self.load_samples(datas)
        # print("加载完成~")

        # 这里做一下处理，将其转换为三元组的形式
        if self.is_train:
            lth = len(datas)//2
            for i in range(lth):
                if datas[i][0] != datas[lth+i][0]:
                    print("error 样本对不匹配")  # 这个信息报了就说明写错了
                self.tripletDatas.append(
                    (datas[i][0], datas[i][1], datas[lth+i][1]))

    #   ind是指在batch中的位置
    def sampels_trans_to_batch(self, lth, att_dimension, batch, ind):
        max_node_size = -float("inf")
        for i in range(lth):
            # max_node_size=max(max_node_size,len(batch[i]["x"]["adj"]))
            max_node_size = max(max_node_size, len(batch[i][ind][0]))

        batch_adj, batch_att, batch_vtype, batch_adj_weight = [], [], [], []
        for i in range(lth):
            temp_adj = np.zeros((max_node_size, max_node_size))
            temp_att = np.zeros((max_node_size, att_dimension))
            temp_vtype = np.zeros((max_node_size, 3))

            sample_size = len(batch[i][ind][0])  # 当前样本的长度

            temp_adj[:sample_size, :sample_size] = batch[i][ind][0]
            temp_att[:sample_size, :] = batch[i][ind][1]
            temp_vtype[:sample_size, :] = batch[i][ind][2]

            batch_adj.append(copy.deepcopy(temp_adj))  # 深拷贝,免得出现错误
            batch_att.append(copy.deepcopy(temp_att))
            batch_vtype.append(copy.deepcopy(temp_vtype))

        res = [torch.tensor(np.array(batch_adj)),
               torch.tensor(np.array(batch_att)),
               torch.tensor(np.array(batch_vtype))]
        return res

    def adjust_samples_to_same_dimension(self, batch):
        lth = len(batch)
        att_dimension = len(batch[0][0][1][0])

        batch_value = []
        res_x, res_y, res_z = [], [], []
        res_x = self.sampels_trans_to_batch(lth, att_dimension, batch, 0)
        res_y = self.sampels_trans_to_batch(lth, att_dimension, batch, 1)
        if self.is_train:
            res_z = self.sampels_trans_to_batch(lth, att_dimension, batch, 2)
        else:
            # value trans to batch
            for i in range(lth):
                batch_value.append(batch[i][3][0])

        return (torch.tensor(batch_value), res_x, res_y, res_z)

    #  其中一个样本
    def __getitem__(self, idx):
        if self.is_train:
            sample_x = self.eSample.get_sample(
                self.tripletDatas[idx][0], modelName)
            sample_y = self.eSample.get_sample(
                self.tripletDatas[idx][1], modelName)
            sample_z = self.eSample.get_sample(
                self.tripletDatas[idx][2], modelName)
            x = [sample_x["adj"], sample_x["att"], sample_x["vtype"]]
            y = [sample_y["adj"], sample_y["att"], sample_y["vtype"]]
            z = [sample_z["adj"], sample_z["att"], sample_z["vtype"]]
            return [x, y, z, [0]]
        else:
            sample_x = self.eSample.get_sample(self.data[idx][0], modelName)
            sample_y = self.eSample.get_sample(self.data[idx][1], modelName)
            x = [sample_x["adj"], sample_x["att"], sample_x["vtype"]]
            y = [sample_y["adj"], sample_y["att"], sample_y["vtype"]]
            value = torch.tensor([self.data[idx][2]])
            return [x, y, [0], value]

    # """
    #     内存放不下,不统一加载了
    # """
    # def load_samples(self,datas):
    #     """
    #         输入为x,y,value
    #         其他的信息不能要哟

    #         增加反汇编实验样本的处理
    #     """
    #     if is_disassemble_experiment==False:
    #         if load_from_shelve and self.datasName in ["train","test","valid"]:
    #             with shelve.open(r"/mnt/mydisk1/chenyongwei/sampleDatas/{}_predata_for_train/datas_{}".format(modelName,self.datasName)) as file:
    #             # with shelve.open(r"/mnt/mydisk1/chenyongwei/malware/BODMAS_dataset/preDataForTrain/{}_{}".format(self.datasName)) as file:
    #                 self.samples=file["samples"]
    #                 file.close()
    #             print("从之前保存的数据中加载")
    #             return
    #         sampleNames = {}
    #         for i in range(len(datas)):
    #             x,y,value=datas[i]
    #             sampleNames[x]=True
    #             sampleNames[y]=True
    #         print("样本对长度:{}, 需加载{}个样本".format(len(datas),len(sampleNames)))
    #         cnt = 0
    #         for sampleName in tqdm(sampleNames.keys()):
    #             sample = self.eSample.get_sample(sampleName,modelName)
    #             self.samples[sampleName]=sample

    #         if  self.datasName in ["train","test","valid"]:
    #             with shelve.open(r"/mnt/mydisk1/chenyongwei/sampleDatas/{}_predata_for_train/datas_{}".format(modelName,self.datasName)) as file:
    #                 file["samples"]=self.samples
    #                 file.close()
    #     else:
    #         # 反汇编实验，加载样本
    #         print("反汇编实验：    正在加载{}处理的样本".format(disassemble_tool))
    #         if load_from_shelve and self.datasName in ["train","test","valid"]:
    #             with shelve.open(r"/home/cyw/projects/function_sim_project/all_data/sampleDatas/functionSim_predata_for_train/{}_datas_{}".format(disassemble_tool,self.datasName)) as file:
    #                 self.samples=file["samples"]
    #                 file.close()
    #             print("从之前保存的数据中加载")
    #             return
    #         for i in tqdm(range(len(datas))):
    #             x,y,value=datas[i]
    #             if x not in self.samples:
    #                 sample = self.eSample.get_sample(x,disassemble_tool)
    #                 self.samples[x]=sample
    #             if y not in self.samples:
    #                 sample = self.eSample.get_sample(y,disassemble_tool)
    #                 self.samples[y]=sample
    #         print("反汇编实验：    {}样本加载完成".format(disassemble_tool))
    #         if self.datasName in ["train","test","valid"]:
    #             with shelve.open(r"/home/cyw/projects/function_sim_project/all_data/sampleDatas/functionSim_predata_for_train/{}_datas_{}".format(disassemble_tool,self.datasName)) as file:
    #                 file["samples"]=self.samples
    #                 file.close()

    def __len__(self):
        if self.is_train:
            return len(self.tripletDatas)
        return len(self.data)


if __name__ == "__main__":
    datapre = dataPre()
    data_pair_infs = datapre.get_pair_infs()
    load_from_shelve = False
    print("处理functionSim数据集的预加载")
    dataset = functionSimDataset(data_pair_infs["test"], "test", False)
    dataset = functionSimDataset(data_pair_infs["valid"], "valid", False)
    dataset = functionSimDataset(data_pair_infs["train"], "train", False)
    print("functionSim数据集的预加载完成")
